﻿using System;
using System.Collections.Generic;
using System.IO;
using System.IO.IsolatedStorage;
using System.Reflection;
using System.Runtime.Serialization.Formatters.Binary;
using System.Text;
using Microsoft.Practices.CompositeUI;
using Microsoft.Practices.CompositeUI.Configuration;
using Microsoft.Practices.CompositeUI.Services;
using Microsoft.Practices.ObjectBuilder;
using Kent.Boogaart.HelperTrinity;
using WPFCABShell.Shell.Interface.Services.Module;
using WPFCABShell.Shell.Interface.Diagnostics;

namespace WPFCABShell.Shell.Services.Module
{
	/// <summary>
	/// An implementation of <see cref="IWhitelistModuleEnumerator"/> for the WPF/CAB shell.
	/// </summary>
	/// <remarks>
	/// <para>
	/// This module enumerator is used to support the plugin infrastructure for the shell. It implicitly trusts
	/// plugins signed with a particular key, and requires other plugins to be explicitly added to a whitelist.
	/// </para>
	/// </remarks>
	public sealed class WhitelistModuleEnumerator : IWhitelistModuleEnumerator, IBuilderAware
	{
		private List<AssemblyName> _whitelist;
		private const string _fileName = "whitelist.bin";
		private static readonly BinaryFormatter _binaryFormatter = new BinaryFormatter();
		private static readonly Log _log = Log.CreateForType(typeof(WhitelistModuleEnumerator));

		//any plugins that have this public key will always be loaded
		//NOTE: you should replace this with the public key from your own .snk
		private static readonly byte[] _trustedPublicKey = new byte[] { 0, 36, 0, 0, 4, 128, 0, 0, 148, 0, 0, 0, 6, 2, 0, 0, 0, 36, 0, 0, 82, 83, 65, 49, 0, 4, 0, 0, 1, 0, 1, 0, 197, 18, 14, 222, 249, 132, 155, 37, 54, 105, 116, 86, 116, 212, 168, 158, 87, 89, 118, 248, 166, 29, 54, 232, 244, 173, 10, 116, 175, 135, 220, 201, 236, 124, 30, 78, 96, 136, 220, 67, 150, 98, 236, 189, 236, 40, 13, 37, 26, 21, 145, 122, 210, 4, 31, 93, 92, 17, 176, 4, 82, 59, 17, 63, 229, 30, 189, 160, 45, 254, 17, 234, 201, 248, 203, 33, 235, 182, 225, 38, 113, 222, 145, 13, 92, 235, 87, 124, 247, 152, 175, 60, 31, 29, 10, 42, 189, 10, 255, 124, 240, 29, 237, 8, 215, 148, 95, 87, 166, 119, 253, 150, 120, 143, 64, 241, 182, 56, 103, 124, 48, 154, 154, 145, 100, 143, 180, 207 };

		private string Path
		{
			get
			{
				return System.IO.Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "Plugins");
			}
		}

		public IEnumerable<IModuleInfo> AllModules
		{
			get
			{
				if (Directory.Exists(Path))
				{
					foreach (string filename in Directory.GetFiles(Path, "*.dll", SearchOption.AllDirectories))
					{
						Assembly assembly = Assembly.ReflectionOnlyLoadFrom(filename);

						if (AssemblyHasModuleAttribute(assembly))
						{
							ModuleInfo moduleInfo = new ModuleInfo();
							moduleInfo.SetAssemblyFile(filename);

							yield return moduleInfo;
						}
					}
				}
			}
		}

		public IModuleInfo[] EnumerateModules()
		{
			_log.Verbose("Enumerating modules:");
			List<ModuleInfo> modules = new List<ModuleInfo>();

			foreach (ModuleInfo moduleInfo in AllModules)
			{
				AssemblyName moduleAssemblyName = GetAssemblyNameForModule(moduleInfo);

				//only enumerate modules on the whitelist
				AssemblyName match = _whitelist.Find(delegate(AssemblyName assemblyName)
				{
					return string.Equals(assemblyName.FullName, moduleAssemblyName.FullName, StringComparison.Ordinal);
				});

				if ((match != null) || IsTrusted(moduleInfo))
				{
					_log.Verbose("Found module: {0}", moduleInfo.AssemblyFile);
					modules.Add(moduleInfo);
				}
			}

			return modules.ToArray();
		}

		public ModuleTrustStatus GetModuleTrustStatus(IModuleInfo moduleInfo)
		{
			ArgumentHelper.AssertNotNull(moduleInfo, "moduleInfo");

			if (IsTrusted(moduleInfo))
			{
				return ModuleTrustStatus.ImplicitlyTrusted;
			}
			else if (IsWhitelisted(moduleInfo))
			{
				return ModuleTrustStatus.Whitelisted;
			}

			return ModuleTrustStatus.NotWhitelisted;
		}

		public ModuleStatus GetStatus(IModuleInfo moduleInfo)
		{
			ArgumentHelper.AssertNotNull(moduleInfo, "moduleInfo");
			AssemblyName moduleAssemblyName = GetAssemblyNameForModule(moduleInfo);

			foreach (Assembly assembly in AppDomain.CurrentDomain.GetAssemblies())
			{
				if (assembly.ReflectionOnly)
				{
					continue;
				}

				if (string.Equals(assembly.FullName, moduleAssemblyName.FullName, StringComparison.Ordinal))
				{
					return ModuleStatus.Loaded;
				}
			}

			return ModuleStatus.NotLoaded;
		}

		public void AddToWhitelist(IModuleInfo moduleInfo)
		{
			ArgumentHelper.AssertNotNull(moduleInfo, "moduleInfo");

			if (IsWhitelisted(moduleInfo))
			{
				return;
			}

			AssemblyName moduleAssemblyName = GetAssemblyNameForModule(moduleInfo);
			_whitelist.Add(moduleAssemblyName);
		}

		public void RemoveFromWhitelist(IModuleInfo moduleInfo)
		{
			ArgumentHelper.AssertNotNull(moduleInfo, "moduleInfo");

			AssemblyName moduleAssemblyName = GetAssemblyNameForModule(moduleInfo);
			_whitelist.RemoveAll(delegate(AssemblyName assemblyName)
			{
				return string.Equals(assemblyName.FullName, moduleAssemblyName.FullName, StringComparison.Ordinal);
			});
		}

		public Assembly GetAssemblyForModule(IModuleInfo moduleInfo)
		{
			ArgumentHelper.AssertNotNull(moduleInfo, "moduleInfo");
			return Assembly.ReflectionOnlyLoadFrom(moduleInfo.AssemblyFile);
		}

		private AssemblyName GetAssemblyNameForModule(IModuleInfo moduleInfo)
		{
			ArgumentHelper.AssertNotNull(moduleInfo, "moduleInfo");
			return AssemblyName.GetAssemblyName(moduleInfo.AssemblyFile);
		}

		private bool IsTrusted(IModuleInfo moduleInfo)
		{
			AssemblyName moduleAssemblyName = GetAssemblyNameForModule(moduleInfo);
			byte[] publicKey = moduleAssemblyName.GetPublicKey();

			if ((publicKey == null) || (publicKey.Length != _trustedPublicKey.Length))
			{
				return false;
			}

			for (uint i = 0; i < publicKey.Length; ++i)
			{
				if (publicKey[i] != _trustedPublicKey[i])
				{
					return false;
				}
			}

			return true;
		}

		private bool IsWhitelisted(IModuleInfo moduleInfo)
		{
			if (IsTrusted(moduleInfo))
			{
				return true;
			}

			AssemblyName moduleAssemblyName = GetAssemblyNameForModule(moduleInfo);

			foreach (AssemblyName assemblyName in _whitelist)
			{
				if (string.Equals(assemblyName.FullName, moduleAssemblyName.FullName, StringComparison.Ordinal))
				{
					return true;
				}
			}

			return false;
		}

		private void Load()
		{
			_log.Verbose("Loading whitelist.");

			try
			{
				IsolatedStorageFile isolatedStorageFile = IsolatedStorageFile.GetUserStoreForAssembly();

				foreach (string fileName in isolatedStorageFile.GetFileNames(_fileName))
				{
					if (string.Equals(fileName, _fileName, StringComparison.OrdinalIgnoreCase))
					{
						//file exists
						using (IsolatedStorageFileStream fileStream = new IsolatedStorageFileStream(_fileName, FileMode.Open, FileAccess.Read, isolatedStorageFile))
						{
							_whitelist = _binaryFormatter.Deserialize(fileStream) as List<AssemblyName>;
						}
					}
				}
			}
			catch (Exception ex)
			{
				_log.Exception(ex);
			}

			if (_whitelist == null)
			{
				_whitelist = new List<AssemblyName>();
			}
		}

		private void Save()
		{
			_log.Verbose("Saving whitelist");

			try
			{
				IsolatedStorageFile isolatedStorageFile = IsolatedStorageFile.GetUserStoreForAssembly();

				using (IsolatedStorageFileStream fileStream = new IsolatedStorageFileStream(_fileName, FileMode.Create, FileAccess.Write, isolatedStorageFile))
				{
					_binaryFormatter.Serialize(fileStream, _whitelist);
				}
			}
			catch (Exception ex)
			{
				_log.Exception(ex);
			}
		}

		void IBuilderAware.OnBuiltUp(string id)
		{
			AppDomain.CurrentDomain.ReflectionOnlyAssemblyResolve += CurrentDomain_ReflectionOnlyAssemblyResolve;
			AppDomain.CurrentDomain.AssemblyResolve += CurrentDomain_AssemblyResolve;
			Load();
		}

		void IBuilderAware.OnTearingDown()
		{
			Save();
			AppDomain.CurrentDomain.ReflectionOnlyAssemblyResolve -= CurrentDomain_ReflectionOnlyAssemblyResolve;
			AppDomain.CurrentDomain.AssemblyResolve -= CurrentDomain_AssemblyResolve;
		}

		private Assembly CurrentDomain_ReflectionOnlyAssemblyResolve(object sender, ResolveEventArgs args)
		{
			return Assembly.ReflectionOnlyLoad(args.Name);
		}

		private Assembly CurrentDomain_AssemblyResolve(object sender, ResolveEventArgs args)
		{
			foreach (Assembly assembly in AppDomain.CurrentDomain.GetAssemblies())
			{
				if (string.Equals(args.Name, assembly.FullName, StringComparison.Ordinal))
				{
					return assembly;
				}
			}

			return null;
		}

		private static bool AssemblyHasModuleAttribute(Assembly assembly)
		{
			IList<CustomAttributeData> customAttributes = CustomAttributeData.GetCustomAttributes(assembly);
			CustomAttributeData moduleAttribute = GetCustomAttributeForType(customAttributes, typeof(ModuleAttribute));
			return moduleAttribute != null;
		}

		private static CustomAttributeData GetCustomAttributeForType(IList<CustomAttributeData> customAttributes, Type type)
		{
			foreach (CustomAttributeData customAttribute in customAttributes)
			{
				if (string.Equals(customAttribute.Constructor.DeclaringType.FullName, type.FullName, StringComparison.Ordinal))
				{
					return customAttribute;
				}
			}

			return null;
		}
	}
}
